#pragma once

#include <ctime>
#include <mutex>
#include <string>
#include <thread>
#include <atomic>
#include <chrono>
#include <sstream>
#include <iomanip>
#include <stdexcept>

#include <mosquitto/mosquitto.h>

namespace mqtt {
class MQTTClient {
public:
    MQTTClient(const std::string &baseClientId = "", bool cleanSession = true)
        : m_mosq(nullptr), m_connected(false), m_shouldReconnect(false),
        m_cleanSession(cleanSession), m_baseClientId(baseClientId)
    {
        mosquitto_lib_init();
        m_clientId = generateClientId();
    }

    ~MQTTClient() {
        disconnect();
        mosquitto_lib_cleanup();
    }

    void disconnect() {
        std::lock_guard<std::mutex> lock(m_mutex);
        m_shouldReconnect = false;

        if (m_mosq) {
            mosquitto_disconnect(m_mosq);
            mosquitto_loop_stop(m_mosq, false);
            mosquitto_destroy(m_mosq);
            m_mosq = nullptr;
        }

        m_connected = false;
    }

    void connect(const std::string &host, uint16_t port, const std::string &username, const std::string &password, uint8_t qos = 0, int keepalive = 60) {
        std::lock_guard<std::mutex> lock(m_mutex);

        m_host = host;
        m_port = port;
        m_username = username;
        m_password = password;
        m_qos = qos;
        m_keepalive = keepalive;

        internalConnect();
        m_shouldReconnect = true;
    }

    void restart(const std::string &host, uint16_t port, const std::string &username, const std::string &password, uint8_t qos) {
        disconnect();
        connect(host, port, username, password, qos);
    }

    int publish(const std::string &topic, const std::string &payload) {
        std::lock_guard<std::mutex> lock(m_mutex);

        if (!m_connected || !m_mosq) {
            std::printf("Not connected to MQTT broker\n");
            return -1;
        }

        int ret = mosquitto_publish(m_mosq, nullptr, topic.c_str(), payload.size(), payload.c_str(), m_qos, false);
        if (ret != MOSQ_ERR_SUCCESS) {
            handleMosquittoError(ret);
        }

        return ret;
    }

    bool isConnected() const {
        return m_connected.load();
    }

private:
    std::string generateClientId() const {
        if (m_baseClientId.empty()) {
            return getFormattedCurrentDateTime();
        }

        return m_baseClientId + "_" + getFormattedCurrentDateTime();
    }

    static std::string getFormattedCurrentDateTime() {
        auto now = std::chrono::system_clock::now();
        auto now_time = std::chrono::system_clock::to_time_t(now);
        auto now_ms = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch());
        auto tm = *std::localtime(&now_time);

        std::stringstream ss;
        char buffer[80];
        std::strftime(buffer, 80, "%Y%m%d%H%M%S", &tm);
        ss << buffer << std::setfill('0') << std::setw(3) << now_ms.count() % 1000;

        return ss.str();
    }

    void internalConnect() {
        if (m_mosq) {
            mosquitto_destroy(m_mosq);
        }

        m_clientId = generateClientId();
        m_mosq = mosquitto_new(m_clientId.empty() ? nullptr : m_clientId.c_str(), m_cleanSession, this);
        if (!m_mosq) {
            throw std::runtime_error("Failed to create mosquitto instance");
        }

        mosquitto_connect_callback_set(m_mosq, onConnect);
        mosquitto_disconnect_callback_set(m_mosq, onDisconnect);

        if (!m_username.empty() && mosquitto_username_pw_set(m_mosq, m_username.c_str(), m_password.c_str()) != MOSQ_ERR_SUCCESS) {
            mosquitto_destroy(m_mosq);
            m_mosq = nullptr;
            throw std::runtime_error("Failed to set username/password");
        }

        int ret = mosquitto_connect_async(m_mosq, m_host.c_str(), m_port, m_keepalive);
        if (ret != MOSQ_ERR_SUCCESS) {
            mosquitto_destroy(m_mosq);
            m_mosq = nullptr;
            handleMosquittoError(ret);
            return;
        }

        ret = mosquitto_loop_start(m_mosq);
        if (ret != MOSQ_ERR_SUCCESS) {
            mosquitto_destroy(m_mosq);
            m_mosq = nullptr;
            handleMosquittoError(ret);
        }
    }

    void handleMosquittoError(int errorCode) {
        switch (errorCode) {
            case MOSQ_ERR_SUCCESS:
                break;

            case MOSQ_ERR_INVAL:
                throw std::runtime_error("Invalid input parameters");

            case MOSQ_ERR_NOMEM:
                throw std::runtime_error("Out of memory");

            case MOSQ_ERR_NO_CONN:
                m_connected = false;
                startReconnect();
                break;

            case MOSQ_ERR_CONN_REFUSED:
                throw std::runtime_error("Connection refused");

            case MOSQ_ERR_CONN_LOST:
                m_connected = false;
                startReconnect();
                break;

            default:
                throw std::runtime_error("MQTT error: " + std::to_string(errorCode));
        }
    }

    void startReconnect() {
        if (!m_shouldReconnect) {
            return;
        }

        std::thread([this]() {
            std::this_thread::sleep_for(std::chrono::seconds(5));

            std::lock_guard<std::mutex> lock(m_mutex);
            if (m_shouldReconnect && !m_connected) {
                try {
                    internalConnect();
                } catch (...) {
                    // If reconnect fails, the next operation will trigger another attempt
                }
            }
        }).detach();
    }

    static void onConnect(struct mosquitto *mosq, void *obj, int rc) {
        MQTTClient *client = static_cast<MQTTClient *>(obj);
        if (rc == MOSQ_ERR_SUCCESS) {
            client->m_connected = true;
        } else {
            client->m_connected = false;
            client->handleMosquittoError(rc);
        }
    }

    static void onDisconnect(struct mosquitto *mosq, void *obj, int rc) {
        MQTTClient *client = static_cast<MQTTClient *>(obj);
        client->m_connected = false;
        client->handleMosquittoError(rc);
    }

private:
    struct mosquitto   *m_mosq;
    std::atomic<bool>  m_connected;
    std::atomic<bool>  m_shouldReconnect;
    bool               m_cleanSession;
    std::string        m_baseClientId;
    std::string        m_clientId;

    std::string        m_host;
    uint16_t           m_port;
    std::string        m_username;
    std::string        m_password;
    uint8_t            m_qos;
    int                m_keepalive;
    mutable std::mutex m_mutex;
};
}
